#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HEX文件加载器和升级流程处理模块

该模块提供以下功能:
1. 解析和加载Intel HEX格式文件
2. 将HEX数据转换为内存中的二进制数据
3. 提供简单的API接口用于访问HEX数据
4. 实现基本的固件升级流程
"""

import os
from typing import Dict, List, Tuple, Optional


class HexLoader:
    """HEX文件加载器类"""
    
    def __init__(self, verify_checksum=True):
        """初始化HEX加载器"""
        self.hex_data = None  # 原始HEX文件内容
        self.memory_map = {}  # 内存映射 {address: data}
        self.segments = []    # 段信息 [(start_address, length), ...]
        self.entry_point = None  # 入口点地址
        self.file_path = None    # 文件路径
        self.file_size = 0       # 文件大小
        self.verify_checksum = verify_checksum  # 是否验证校验和
        
    def load_hex_file(self, file_path: str) -> bool:
        """
        加载HEX文件到内存
        
        Args:
            file_path (str): HEX文件路径
            
        Returns:
            bool: 加载成功返回True，否则返回False
        """
        try:
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"文件不存在: {file_path}")
            
            self.file_path = file_path
            self.file_size = os.path.getsize(file_path)
            
            # 尝试不同的编码方式读取HEX文件
            encodings = ['utf-8-sig', 'utf-16', 'utf-16-le', 'utf-16-be', 'gbk', 'latin1']
            self.hex_data = None
            
            for encoding in encodings:
                try:
                    with open(file_path, 'r', encoding=encoding) as file:
                        self.hex_data = file.read()
                    break
                except UnicodeDecodeError:
                    continue
            
            if self.hex_data is None:
                # 如果所有编码都失败，以二进制模式读取并尝试解码
                with open(file_path, 'rb') as file:
                    raw_data = file.read()
                # 尝试解码为文本
                self.hex_data = raw_data.decode('utf-8-sig', errors='ignore')
            
            # 解析HEX数据
            return self._parse_hex_data()
            
        except Exception as e:
            print(f"加载HEX文件失败: {str(e)}")
            return False
    
    def _parse_hex_data(self) -> bool:
        """
        解析HEX数据并构建内存映射
        
        Returns:
            bool: 解析成功返回True，否则返回False
        """
        try:
            lines = self.hex_data.strip().split('\n')
            self.memory_map = {}
            self.segments = []
            
            current_address = 0
            segment_start = None
            segment_length = 0
            
            for line in lines:
                if not line.startswith(':'):
                    continue
                
                # 解析HEX行数据
                data_length = int(line[1:3], 16)
                address = int(line[3:7], 16)
                record_type = int(line[7:9], 16)
                data = line[9:9 + data_length * 2]
                checksum = int(line[9 + data_length * 2:11 + data_length * 2], 16)
                
                # 验证校验和（如果启用）
                if self.verify_checksum and not self._verify_checksum(line):
                    raise ValueError(f"HEX行校验和错误: {line}")
                
                # 处理不同类型的记录
                if record_type == 0:  # 数据记录
                    # 计算实际地址
                    actual_address = current_address + address
                    
                    # 存储数据到内存映射
                    for i in range(0, len(data), 2):
                        byte_address = actual_address + (i // 2)
                        byte_value = int(data[i:i+2], 16)
                        self.memory_map[byte_address] = byte_value
                    
                    # 更新段信息
                    if segment_start is None:
                        segment_start = actual_address
                    segment_length += data_length
                    
                elif record_type == 1:  # 结束记录
                    if segment_start is not None:
                        self.segments.append((segment_start, segment_length))
                    break
                    
                elif record_type == 2:  # 扩展段地址记录
                    esa = int(data, 16)
                    current_address = esa * 16
                    
                elif record_type == 3:  # 起始段地址记录
                    pass  # 通常用于16位处理器
                    
                elif record_type == 4:  # 扩展线性地址记录
                    ela = int(data, 16)
                    current_address = ela * 65536
                    
                elif record_type == 5:  # 起始线性地址记录
                    self.entry_point = int(data, 16)
                    
            # 添加最后一个段
            if segment_start is not None and segment_length > 0:
                self.segments.append((segment_start, segment_length))
                
            return True
            
        except Exception as e:
            print(f"解析HEX数据失败: {str(e)}")
            return False
    
    def _verify_checksum(self, line: str) -> bool:
        """
        验证HEX行的校验和
        
        Args:
            line (str): HEX行数据
            
        Returns:
            bool: 校验和正确返回True，否则返回False
        """
        try:
            # 提取所有字节（包括校验和）
            data_bytes = []
            for i in range(1, len(line), 2):
                data_bytes.append(int(line[i:i+2], 16))
            
            # 计算校验和，所有字节之和的低8位应该为0
            checksum = sum(data_bytes) & 0xFF
            
            return checksum == 0
            
        except Exception:
            return False
    
    def get_memory_data(self, address: int, length: int = 1) -> Optional[bytes]:
        """
        从内存映射中获取数据
        
        Args:
            address (int): 起始地址
            length (int): 数据长度
            
        Returns:
            Optional[bytes]: 获取的数据，失败返回None
        """
        try:
            data = bytearray()
            for i in range(length):
                addr = address + i
                if addr in self.memory_map:
                    data.append(self.memory_map[addr])
                else:
                    data.append(0)  # 未定义的内存区域填充0
            return bytes(data)
        except Exception:
            return None
    
    def get_segment_data(self, segment_index: int = 0) -> Optional[Tuple[int, bytes]]:
        """
        获取指定段的数据
        
        Args:
            segment_index (int): 段索引
            
        Returns:
            Optional[Tuple[int, bytes]]: (起始地址, 数据)，失败返回None
        """
        try:
            if segment_index >= len(self.segments):
                return None
            
            start_address, length = self.segments[segment_index]
            data = self.get_memory_data(start_address, length)
            if data is not None:
                return (start_address, data)
            return None
        except Exception:
            return None
    
    def get_file_info(self) -> Dict:
        """
        获取文件信息
        
        Returns:
            Dict: 文件信息字典
        """
        return {
            'file_path': self.file_path,
            'file_size': self.file_size,
            'segments_count': len(self.segments),
            'total_bytes': len(self.memory_map),
            'entry_point': self.entry_point,
            'segments': self.segments
        }
    
    def get_memory_map_info(self) -> Dict:
        """
        获取内存映射信息
        
        Returns:
            Dict: 内存映射信息
        """
        if not self.memory_map:
            return {}
        
        addresses = list(self.memory_map.keys())
        addresses.sort()
        
        return {
            'min_address': min(addresses) if addresses else 0,
            'max_address': max(addresses) if addresses else 0,
            'address_range': max(addresses) - min(addresses) + 1 if addresses else 0,
            'used_addresses': len(addresses)
        }
    
    def export_to_bin(self, output_path: str) -> bool:
        """
        将HEX数据导出为BIN文件
        
        Args:
            output_path (str): 输出文件路径
            
        Returns:
            bool: 导出成功返回True，否则返回False
        """
        try:
            if not self.memory_map:
                raise ValueError("没有加载的HEX数据")
            
            # 获取地址范围
            addresses = list(self.memory_map.keys())
            min_addr = min(addresses)
            max_addr = max(addresses)
            
            # 创建连续的二进制数据
            bin_data = bytearray(max_addr - min_addr + 1)
            for addr, value in self.memory_map.items():
                bin_data[addr - min_addr] = value
            
            # 写入文件
            with open(output_path, 'wb') as f:
                f.write(bin_data)
            
            return True
        except Exception as e:
            print(f"导出BIN文件失败: {str(e)}")
            return False


class FirmwareUpdater:
    """固件升级器类"""
    
    def __init__(self, hex_loader: HexLoader):
        """
        初始化固件升级器
        
        Args:
            hex_loader (HexLoader): HEX加载器实例
        """
        self.hex_loader = hex_loader
        self.upgrade_steps = []  # 升级步骤
        self.current_step = 0    # 当前步骤
        self.is_upgrading = False  # 是否正在升级
        
    def prepare_upgrade(self) -> bool:
        """
        准备升级流程
        
        Returns:
            bool: 准备成功返回True，否则返回False
        """
        try:
            if not self.hex_loader.memory_map:
                raise ValueError("没有加载的HEX数据")
            
            # 清空之前的升级步骤
            self.upgrade_steps = []
            self.current_step = 0
            
            # 添加默认升级步骤
            self._add_default_upgrade_steps()
            
            return True
        except Exception as e:
            print(f"准备升级失败: {str(e)}")
            return False
    
    def _add_default_upgrade_steps(self):
        """添加默认升级步骤"""
        self.upgrade_steps = [
            "连接设备",
            "擦除Flash",
            "写入固件数据",
            "验证固件数据",
            "重启设备"
        ]
    
    def add_upgrade_step(self, step: str):
        """
        添加自定义升级步骤
        
        Args:
            step (str): 升级步骤描述
        """
        self.upgrade_steps.append(step)
    
    def get_upgrade_steps(self) -> List[str]:
        """
        获取所有升级步骤
        
        Returns:
            List[str]: 升级步骤列表
        """
        return self.upgrade_steps.copy()
    
    def start_upgrade(self) -> bool:
        """
        开始升级流程
        
        Returns:
            bool: 升级开始成功返回True，否则返回False
        """
        try:
            if not self.upgrade_steps:
                raise ValueError("没有定义升级步骤")
            
            self.is_upgrading = True
            self.current_step = 0
            return True
        except Exception as e:
            print(f"开始升级失败: {str(e)}")
            return False
    
    def next_step(self) -> Tuple[bool, Optional[str]]:
        """
        执行下一步升级
        
        Returns:
            Tuple[bool, Optional[str]]: (是否还有下一步, 当前步骤描述)
        """
        if not self.is_upgrading:
            return (False, None)
        
        if self.current_step >= len(self.upgrade_steps):
            self.is_upgrading = False
            return (False, None)
        
        current_step_desc = self.upgrade_steps[self.current_step]
        self.current_step += 1
        
        # 检查是否还有下一步
        has_next = self.current_step < len(self.upgrade_steps)
        
        return (has_next, current_step_desc)
    
    def get_progress(self) -> Tuple[int, int]:
        """
        获取升级进度
        
        Returns:
            Tuple[int, int]: (当前步骤, 总步骤数)
        """
        return (self.current_step, len(self.upgrade_steps))
    
    def cancel_upgrade(self):
        """取消升级"""
        self.is_upgrading = False
        self.current_step = 0


# 使用示例和测试函数
def main():
    """主函数，用于测试和演示"""
    print("HEX加载器和升级流程演示")
    
    # 创建HEX加载器实例（不验证校验和）
    loader = HexLoader(verify_checksum=False)
    
    # 加载测试HEX文件
    if loader.load_hex_file("test.hex"):
        print("HEX文件加载成功")
        
        # 显示文件信息
        info = loader.get_file_info()
        print(f"文件路径: {info['file_path']}")
        print(f"文件大小: {info['file_size']} 字节")
        print(f"段数量: {info['segments_count']}")
        print(f"总字节数: {info['total_bytes']}")
        print(f"入口点: 0x{info['entry_point']:08X}" if info['entry_point'] else "入口点: 无")
        print(f"段信息: {info['segments']}")
        
        # 显示内存映射信息
        mem_info = loader.get_memory_map_info()
        print(f"最小地址: 0x{mem_info['min_address']:08X}")
        print(f"最大地址: 0x{mem_info['max_address']:08X}")
        print(f"地址范围: {mem_info['address_range']} 字节")
        print(f"使用地址数: {mem_info['used_addresses']}")
        
        # 创建升级器实例
        updater = FirmwareUpdater(loader)
        
        # 准备升级
        if updater.prepare_upgrade():
            print("\n升级准备完成")
            
            # 显示升级步骤
            steps = updater.get_upgrade_steps()
            print("升级步骤:")
            for i, step in enumerate(steps, 1):
                print(f"  {i}. {step}")
            
            # 模拟升级过程
            print("\n开始升级...")
            if updater.start_upgrade():
                while True:
                    has_next, step_desc = updater.next_step()
                    if step_desc:
                        progress = updater.get_progress()
                        print(f"执行步骤: {step_desc} ({progress[0]}/{progress[1]})")
                    
                    if not has_next:
                        break
                
                print("升级完成!")
            else:
                print("升级启动失败")
        else:
            print("升级准备失败")
    else:
        print("HEX文件加载失败")


if __name__ == "__main__":
    main()
